{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import KFold, StratifiedKFold\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from sklearn.linear_model import LogisticRegression, LinearRegression\n",
    "\n",
    "from econml.sklearn_extensions.linear_model import StatsModelsLinearRegression\n",
    "from econml.utilities import _safe_norm_ppf\n",
    "\n",
    "\n",
    "def dr_sim(model_t, model_y, n_folds, Y, T, X, min_t=0.01):\n",
    "    \"\"\"\n",
    "    Doubly-robust learner\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model_t : object\n",
    "        A classifier for the treatment.\n",
    "    model_y : object\n",
    "        A regressor for the outcome.\n",
    "    n_folds : int\n",
    "        The number of folds to use in cross-validation.\n",
    "    Y : array-like\n",
    "        The outcome variable.\n",
    "    T : array-like\n",
    "        The treatment variable.\n",
    "    X : array-like\n",
    "        The covariates.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    theta : array-like\n",
    "        The estimated treatment effect per treatment.\n",
    "    sigma : array-like\n",
    "        The estimated variance of the outcome residuals per treatment.\n",
    "    nu : array-like\n",
    "        The estimated variance of the treatment per treatment.\n",
    "    cov : array-like\n",
    "        The covariance matrix of the estimated parameters, of shape (n_treatments, 3, 3).\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize KFold cross-validation\n",
    "    kf = StratifiedKFold(n_splits=n_folds, shuffle=True)\n",
    "    T_ohe = OneHotEncoder(sparse_output=False, drop='first').fit_transform(T.reshape(-1, 1))\n",
    "\n",
    "    # Initialize arrays to hold predictions\n",
    "    y_pred = np.zeros((Y.shape[0], T_ohe.shape[1]+1)) # include prediction for T=0\n",
    "    t_pred = np.zeros((T_ohe.shape[0], T_ohe.shape[1]+1)) # include prediction for T=0\n",
    "    alpha = np.zeros_like(T_ohe)\n",
    "\n",
    "    # one value of theta, sigma, nu for each non-control treatment\n",
    "    theta = np.zeros(T_ohe.shape[1])\n",
    "    sigma = np.zeros(T_ohe.shape[1])\n",
    "    nu = np.zeros(T_ohe.shape[1])\n",
    "\n",
    "    # one theta, sigma, nu covariance matrix for each non-control treatment\n",
    "    cov = np.zeros((T_ohe.shape[1], 3, 3))\n",
    "\n",
    "    # Cross-validation loop\n",
    "    for train_index, test_index in kf.split(X,T):\n",
    "        X_train, X_test = X[train_index], X[test_index]\n",
    "        Y_train, Y_test = Y[train_index], Y[test_index]\n",
    "        T_train = T[train_index]\n",
    "        To_train, To_test = T_ohe[train_index], T_ohe[test_index]\n",
    "     \n",
    "        # Fit the treatment model\n",
    "        t_pred[test_index] = model_t.fit(X_train, T_train).predict_proba(X_test)\n",
    "\n",
    "        t_pred[test_index] = np.clip(t_pred[test_index], min_t, 1-min_t)  # avoid division by zero\n",
    "        t_pred[test_index] = t_pred[test_index] / np.sum(t_pred[test_index], axis=1, keepdims=True)  # normalize to sum to 1\n",
    "        \n",
    "        # Fit the outcome model\n",
    "        model_y.fit(np.hstack([X_train, To_train]), Y_train)\n",
    "        T_hypo = np.zeros_like(To_test)\n",
    "        y_pred[test_index,0] += model_y.predict(np.hstack([X_test, T_hypo]))\n",
    "        for i in range(To_train.shape[1]):\n",
    "            T_hypo = np.zeros_like(To_test)\n",
    "            T_hypo[:,i] = 1\n",
    "            y_pred[test_index,i+1] += model_y.predict(np.hstack([X_test, T_hypo]))\n",
    "\n",
    "    # ATE, sigma^2, and nu^2\n",
    "    for i in range(T_ohe.shape[1]):\n",
    "        theta_score = y_pred[:,i+1] - y_pred[:,0] + (Y-y_pred[:,i+1]) * (T_ohe[:,i] == 1)/t_pred[:,i+1] - (Y-y_pred[:,0]) * (np.all(T_ohe==0, axis=1)/t_pred[:,0])\n",
    "        sigma_score = (Y-np.choose(T, y_pred.T))**2 # exclude rows with other treatments\n",
    "        alpha[:,i] = (T_ohe[:,i] == 1)/t_pred[:,i+1] - (np.all(T_ohe==0, axis=1))/t_pred[:,0]\n",
    "        nu_score = 2*(1/t_pred[:,i+1]+1/t_pred[:,0])-alpha[:,i]**2\n",
    "        theta[i] = np.mean(theta_score)\n",
    "        sigma[i] = np.mean(sigma_score)\n",
    "        nu[i] = np.mean(nu_score)\n",
    "        scores = np.stack([theta_score-theta[i], sigma_score-sigma[i], nu_score-nu[i]], axis=1)\n",
    "        cov[i,:,:] = (scores.T @ scores / len(scores)) / len(scores)\n",
    "\n",
    "    return theta, sigma, nu, cov\n",
    "\n",
    "def sensitivity_interval(theta, sigma, nu, cov, alpha, c_y, c_t, rho):\n",
    "    \"\"\"\n",
    "    Calculate the sensitivity interval for a doubly-robust learner.\n",
    "    \"\"\"\n",
    "    C = np.abs(rho) * np.sqrt(c_y) * np.sqrt(c_t/(1-c_t))/2\n",
    "    ests = np.array([theta, sigma, nu])\n",
    "    \n",
    "    coefs_p = np.array([1, C*np.sqrt(nu/sigma), C*np.sqrt(sigma/nu)])\n",
    "    coefs_n = np.array([1, -C*np.sqrt(nu/sigma), -C*np.sqrt(sigma/nu)])\n",
    "    # One dimensional normal distribution:\n",
    "    sigma_p = coefs_p @ cov @ coefs_p\n",
    "    sigma_n = coefs_n @ cov @ coefs_n\n",
    "\n",
    "    # print(f\"theta bounds: {ests @ coefs_n}, {ests @ coefs_p}\")\n",
    "    # print(f\"sigma bounds: {sigma_n}, {sigma_p}\")\n",
    "\n",
    "    lb = _safe_norm_ppf(alpha / 2, loc=ests @ coefs_n, scale=np.sqrt(sigma_n))\n",
    "    ub = _safe_norm_ppf(1 - alpha / 2, loc=ests @ coefs_p, scale=np.sqrt(sigma_p))\n",
    "\n",
    "    return (lb, ub)\n",
    "\n",
    "\n",
    "def RV(theta, sigma, nu, cov, alpha):\n",
    "    # The robustness value is the degree of confounding of *both* the treatment and the outcome that still produces an interval\n",
    "    # that excludes zero.\n",
    "\n",
    "    # We're looking for a value of r such that the sensitivity bounds just touch zero\n",
    "\n",
    "    r = 0\n",
    "    r_up = 1\n",
    "    r_down = 0\n",
    "    lb, ub = sensitivity_interval(theta, sigma, nu, cov, alpha, 0, 0, 1)\n",
    "    if lb < 0 and ub > 0:\n",
    "        return 0\n",
    "    \n",
    "    else:\n",
    "        if lb > 0:\n",
    "            target = 0\n",
    "            mult = 1\n",
    "            d = lb\n",
    "        else:\n",
    "            target = 1\n",
    "            mult = -1\n",
    "            d = ub\n",
    "\n",
    "    while abs(d) > 1e-6:\n",
    "        d = mult * sensitivity_interval(theta, sigma, nu, cov, alpha, r, r, 1)[target]\n",
    "        if d > 0:\n",
    "            r_down = r\n",
    "        else:\n",
    "            r_up = r\n",
    "\n",
    "        r = (r_down + r_up) / 2\n",
    "        \n",
    "    return r\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulate data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(3)\n",
    "n = 10000\n",
    "alpha = np.random.normal(size=5)\n",
    "X = np.random.normal(size=(n,5), loc=[1,0.5,0,0,0])/3 # closest to center 1, then 2, then 3\n",
    "centers = np.array([[1,0,0,0,0], [0,1,0,0,0], [0,0,1,0,0]]) # trinary treatment\n",
    "# centers = np.array([[1,0,0,0,0], [0,1,0,0,0]]) # uncomment for binary treatment\n",
    "\n",
    "ds = X[:,None,:]-centers[None,:,:]\n",
    "ds = np.einsum(\"nci,nci->nc\", ds, ds)\n",
    "\n",
    "ps_r = np.exp(-ds)\n",
    "ps = ps_r / np.sum(ps_r, axis=1, keepdims=True)\n",
    "\n",
    "T = np.random.default_rng().multinomial(1, ps) @ np.arange(len(centers))\n",
    "\n",
    "Y = np.random.normal(size=n) + 3*(T == 1)*X[:,1] - (T == 2) + 2 * X @ alpha\n",
    "\n",
    "true_theta = np.mean(3*X[:,1]), -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run simple dml and calculate intermediate values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig_level = 0.1\n",
    "\n",
    "theta, sigma, nu, sig = dr_sim(LogisticRegression(), LinearRegression(), 2, Y, T, X)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculate a \"sensitivity interval\". \n",
    "\n",
    "Need to supply \"strength of latent confounder\" as argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'i': 0, 'lb': -2.4222316375247765, 'ub': 3.4062936431220736}\n",
      "{'i': 1, 'lb': -4.248998723655173, 'ub': 2.1903767876822107}\n"
     ]
    }
   ],
   "source": [
    "for i in range(theta.shape[0]):\n",
    "    lb, ub = sensitivity_interval(theta[i], sigma[i], nu[i], sig[i], sig_level, 0.6, 0.6, 1)\n",
    "    print({'i': i, 'lb': lb, 'ub': ub})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculate Robustness Value. \n",
    "\n",
    "The required strength of a latent confounder in order for the confidence interval to include 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'i': 0, 'RV': 0.1342604160308838}\n",
      "{'i': 1, 'RV': 0.25589704513549805}\n"
     ]
    }
   ],
   "source": [
    "for i in range(theta.shape[0]):\n",
    "    rv = RV(theta[i], sigma[i], nu[i], sig[i], sig_level)\n",
    "    print({'i': i, 'RV': rv})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ablations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'i': 1,\n",
       "  'alpha': 0.1,\n",
       "  't_ind': 1,\n",
       "  'sensitivity_interval': (-5.595590315127781, 3.648832519620443),\n",
       "  'RV': 0.17194628715515137},\n",
       " {'i': 5,\n",
       "  'alpha': 0.1,\n",
       "  't_ind': 1,\n",
       "  'sensitivity_interval': (-4.242302203189104, 2.1724266319271592),\n",
       "  'RV': 0.2583136558532715},\n",
       " {'i': 15,\n",
       "  'alpha': 0.1,\n",
       "  't_ind': 1,\n",
       "  'sensitivity_interval': (-4.252962106477528, 2.1979664010821613),\n",
       "  'RV': 0.2550685405731201},\n",
       " {'i': 30,\n",
       "  'alpha': 0.1,\n",
       "  't_ind': 1,\n",
       "  'sensitivity_interval': (-4.247614757215935, 2.1916647643905867),\n",
       "  'RV': 0.2557647228240967}]"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_level = 0.1\n",
    "t_ind = 1\n",
    "\n",
    "results = []\n",
    "for i in [1, 5, 15, 30]:\n",
    "    theta, sigma, nu, sig = dr_sim(LogisticRegression(), LinearRegression(), 2, Y, T, X[:,:i])\n",
    "    result_dict = {\n",
    "        'i': i,\n",
    "        'alpha': sig_level,\n",
    "        't_ind': t_ind,\n",
    "        'sensitivity_interval': sensitivity_interval(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level, 0.6, 0.6, 1),\n",
    "        'RV': RV(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level)\n",
    "    }\n",
    "    results.append(result_dict)\n",
    "                   \n",
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>i</th>\n",
       "      <th>alpha</th>\n",
       "      <th>t_ind</th>\n",
       "      <th>sensitivity_interval</th>\n",
       "      <th>RV</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>1</td>\n",
       "      <td>(-5.595590315127781, 3.648832519620443)</td>\n",
       "      <td>0.171946</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5</td>\n",
       "      <td>0.1</td>\n",
       "      <td>1</td>\n",
       "      <td>(-4.242302203189104, 2.1724266319271592)</td>\n",
       "      <td>0.258314</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15</td>\n",
       "      <td>0.1</td>\n",
       "      <td>1</td>\n",
       "      <td>(-4.252962106477528, 2.1979664010821613)</td>\n",
       "      <td>0.255069</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>30</td>\n",
       "      <td>0.1</td>\n",
       "      <td>1</td>\n",
       "      <td>(-4.247614757215935, 2.1916647643905867)</td>\n",
       "      <td>0.255765</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    i  alpha  t_ind                      sensitivity_interval        RV\n",
       "0   1    0.1      1   (-5.595590315127781, 3.648832519620443)  0.171946\n",
       "1   5    0.1      1  (-4.242302203189104, 2.1724266319271592)  0.258314\n",
       "2  15    0.1      1  (-4.252962106477528, 2.1979664010821613)  0.255069\n",
       "3  30    0.1      1  (-4.247614757215935, 2.1916647643905867)  0.255765"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    pd.DataFrame(results)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DoubleML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "import doubleml as dml\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Y</th>\n",
       "      <th>T</th>\n",
       "      <th>X0</th>\n",
       "      <th>X1</th>\n",
       "      <th>X2</th>\n",
       "      <th>X3</th>\n",
       "      <th>X4</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.025162</td>\n",
       "      <td>1</td>\n",
       "      <td>0.215080</td>\n",
       "      <td>0.139086</td>\n",
       "      <td>-0.209000</td>\n",
       "      <td>-0.014606</td>\n",
       "      <td>-0.159073</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.576181</td>\n",
       "      <td>1</td>\n",
       "      <td>-0.104622</td>\n",
       "      <td>0.461541</td>\n",
       "      <td>0.293773</td>\n",
       "      <td>0.569858</td>\n",
       "      <td>0.016678</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-1.676531</td>\n",
       "      <td>1</td>\n",
       "      <td>0.198441</td>\n",
       "      <td>-0.015120</td>\n",
       "      <td>-0.515492</td>\n",
       "      <td>0.327456</td>\n",
       "      <td>-0.367023</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.059101</td>\n",
       "      <td>1</td>\n",
       "      <td>-0.061682</td>\n",
       "      <td>0.098117</td>\n",
       "      <td>0.495383</td>\n",
       "      <td>0.078905</td>\n",
       "      <td>-0.341262</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.764248</td>\n",
       "      <td>1</td>\n",
       "      <td>0.095669</td>\n",
       "      <td>0.375082</td>\n",
       "      <td>-0.053504</td>\n",
       "      <td>-0.256279</td>\n",
       "      <td>-0.076677</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          Y  T        X0        X1        X2        X3        X4\n",
       "0  1.025162  1  0.215080  0.139086 -0.209000 -0.014606 -0.159073\n",
       "1  0.576181  1 -0.104622  0.461541  0.293773  0.569858  0.016678\n",
       "2 -1.676531  1  0.198441 -0.015120 -0.515492  0.327456 -0.367023\n",
       "3  0.059101  1 -0.061682  0.098117  0.495383  0.078905 -0.341262\n",
       "4  2.764248  1  0.095669  0.375082 -0.053504 -0.256279 -0.076677"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.concat([\n",
    "    pd.DataFrame(Y).squeeze().to_frame('Y'),\n",
    "    pd.DataFrame(T).squeeze().to_frame('T'),\n",
    "    pd.DataFrame(X).add_prefix('X'),\n",
    "], axis=1)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<doubleml.double_ml_data.DoubleMLData at 0x1f0917cfca0>"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dml_data = dml.DoubleMLData(df, 'Y', 'T')\n",
    "dml_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Incompatible data. To fit an IRM model with DML exactly one binary variable with values 0 and 1 needs to be specified as treatment variable.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[96], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m dml_obj \u001b[38;5;241m=\u001b[39m \u001b[43mdml\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mDoubleMLIRM\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdml_data\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m      2\u001b[0m \u001b[43m                          \u001b[49m\u001b[43mml_g\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mLinearRegression\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m      3\u001b[0m \u001b[43m                          \u001b[49m\u001b[43mml_m\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mLogisticRegression\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m      4\u001b[0m \u001b[43m                          \u001b[49m\u001b[43mn_folds\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m      5\u001b[0m \u001b[43m                          \u001b[49m\u001b[43mscore\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mpartialling out\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      6\u001b[0m dml_obj\u001b[38;5;241m.\u001b[39mfit()\n\u001b[0;32m      7\u001b[0m \u001b[38;5;28mprint\u001b[39m(dml_obj)\n",
      "File \u001b[1;32mc:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\irm\\irm.py:134\u001b[0m, in \u001b[0;36mDoubleMLIRM.__init__\u001b[1;34m(self, obj_dml_data, ml_g, ml_m, n_folds, n_rep, score, weights, normalize_ipw, trimming_rule, trimming_threshold, draw_sample_splitting)\u001b[0m\n\u001b[0;32m    116\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m    117\u001b[0m              obj_dml_data,\n\u001b[0;32m    118\u001b[0m              ml_g,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    126\u001b[0m              trimming_threshold\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e-2\u001b[39m,\n\u001b[0;32m    127\u001b[0m              draw_sample_splitting\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m    128\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(obj_dml_data,\n\u001b[0;32m    129\u001b[0m                      n_folds,\n\u001b[0;32m    130\u001b[0m                      n_rep,\n\u001b[0;32m    131\u001b[0m                      score,\n\u001b[0;32m    132\u001b[0m                      draw_sample_splitting)\n\u001b[1;32m--> 134\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_check_data\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dml_data\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    135\u001b[0m     valid_scores \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mATE\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mATTE\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[0;32m    136\u001b[0m     _check_score(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mscore, valid_scores, allow_callable\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "File \u001b[1;32mc:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\irm\\irm.py:247\u001b[0m, in \u001b[0;36mDoubleMLIRM._check_data\u001b[1;34m(self, obj_dml_data)\u001b[0m\n\u001b[0;32m    245\u001b[0m zero_one_treat \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mall((np\u001b[38;5;241m.\u001b[39mpower(obj_dml_data\u001b[38;5;241m.\u001b[39md, \u001b[38;5;241m2\u001b[39m) \u001b[38;5;241m-\u001b[39m obj_dml_data\u001b[38;5;241m.\u001b[39md) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m    246\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (one_treat \u001b[38;5;241m&\u001b[39m binary_treat \u001b[38;5;241m&\u001b[39m zero_one_treat):\n\u001b[1;32m--> 247\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mIncompatible data. \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    248\u001b[0m                      \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mTo fit an IRM model with DML \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    249\u001b[0m                      \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mexactly one binary variable with values 0 and 1 \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    250\u001b[0m                      \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mneeds to be specified as treatment variable.\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m    251\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m\n",
      "\u001b[1;31mValueError\u001b[0m: Incompatible data. To fit an IRM model with DML exactly one binary variable with values 0 and 1 needs to be specified as treatment variable."
     ]
    }
   ],
   "source": [
    "dml_obj = dml.DoubleMLIRM(dml_data,\n",
    "                          ml_g=LinearRegression(),\n",
    "                          ml_m=LogisticRegression(),\n",
    "                          n_folds=2,\n",
    "                          score='partialling out',)\n",
    "dml_obj.fit()\n",
    "print(dml_obj)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== Sensitivity Analysis ==================\n",
      "\n",
      "------------------ Scenario          ------------------\n",
      "Significance Level: level=0.95\n",
      "Sensitivity parameters: cf_y=0.03; cf_d=0.03, rho=1.0\n",
      "\n",
      "------------------ Bounds with CI    ------------------\n",
      "   CI lower  theta lower     theta  theta upper  CI upper\n",
      "0 -0.471507       -0.446 -0.397119    -0.348237 -0.322857\n",
      "\n",
      "------------------ Robustness Values ------------------\n",
      "   H_0     RV (%)    RVa (%)\n",
      "0  0.0  21.873191  20.629383\n"
     ]
    }
   ],
   "source": [
    "dml_obj.sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.)\n",
    "print(dml_obj.sensitivity_summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       cf_y      cf_d  rho  delta_theta\n",
      "T  0.019825  0.000095 -1.0      -0.0022\n"
     ]
    }
   ],
   "source": [
    "sens_benchmark = dml_obj.sensitivity_benchmark(benchmarking_set=[\"X4\"])\n",
    "print(sens_benchmark)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# New datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DoubleML confounded synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from doubleml.datasets import make_confounded_plr_data, make_confounded_irm_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\datasets.py:1050: UserWarning: Propensity score is close to 0 or 1. Trimming is at 0.01 and 0.99 is applied\n",
      "  warnings.warn(f'Propensity score is close to 0 or 1. '\n"
     ]
    }
   ],
   "source": [
    "cf_y = 0.1\n",
    "cf_d = 0.1\n",
    "theta = 5.0\n",
    "dpg_dict = make_confounded_irm_data(n_obs=10000, cf_y=cf_y, cf_d=cf_d, theta=theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_cols = [f'X{i + 1}' for i in np.arange(dpg_dict['x'].shape[1])]\n",
    "df = pd.DataFrame(np.column_stack((dpg_dict['x'], dpg_dict['y'], dpg_dict['d'])), columns=x_cols + ['y', 'd'])\n",
    "dml_data = dml.DoubleMLData(df, 'y', 'd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== DoubleMLIRM Object ==================\n",
      "\n",
      "------------------ Data summary      ------------------\n",
      "Outcome variable: y\n",
      "Treatment variable(s): ['d']\n",
      "Covariates: ['X1', 'X2', 'X3', 'X4', 'X5']\n",
      "Instrument variable(s): None\n",
      "No. Observations: 10000\n",
      "\n",
      "------------------ Score & algorithm ------------------\n",
      "Score function: ATE\n",
      "\n",
      "------------------ Machine learner   ------------------\n",
      "Learner ml_g: RandomForestRegressor()\n",
      "Learner ml_m: RandomForestClassifier()\n",
      "Out-of-sample Performance:\n",
      "Regression:\n",
      "Learner ml_g0 RMSE: [[1.08528722]]\n",
      "Learner ml_g1 RMSE: [[1.14348731]]\n",
      "Classification:\n",
      "Learner ml_m Log Loss: [[0.67196054]]\n",
      "\n",
      "------------------ Resampling        ------------------\n",
      "No. folds: 2\n",
      "No. repeated sample splits: 1\n",
      "\n",
      "------------------ Fit summary       ------------------\n",
      "       coef   std err          t  P>|t|   2.5 %   97.5 %\n",
      "d  5.145625  0.062208  82.716619    0.0  5.0237  5.26755\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\utils\\_checks.py:205: UserWarning: Propensity predictions from learner RandomForestClassifier() for ml_m are close to zero or one (eps=1e-12).\n",
      "  warnings.warn(f'Propensity predictions from learner {str(learner)} for'\n"
     ]
    }
   ],
   "source": [
    "dml_obj = dml.DoubleMLIRM(dml_data,\n",
    "                          ml_g=RandomForestRegressor(),\n",
    "                          ml_m=RandomForestClassifier(),\n",
    "                          n_folds=2,\n",
    "                          score='ATE',)\n",
    "dml_obj.fit()\n",
    "print(dml_obj)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== Sensitivity Analysis ==================\n",
      "\n",
      "------------------ Scenario          ------------------\n",
      "Significance Level: level=0.95\n",
      "Sensitivity parameters: cf_y=0.1; cf_d=0.1, rho=1.0\n",
      "\n",
      "------------------ Bounds with CI    ------------------\n",
      "   CI lower  theta lower     theta  theta upper  CI upper\n",
      "0  4.858139     4.961768  5.145625     5.329482  5.466168\n",
      "\n",
      "------------------ Robustness Values ------------------\n",
      "   H_0     RV (%)    RVa (%)\n",
      "0  0.0  90.573845  84.951031\n"
     ]
    }
   ],
   "source": [
    "dml_obj.sensitivity_analysis(cf_y=0.1, cf_d=0.1, rho=1.)\n",
    "print(dml_obj.sensitivity_summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'alpha': 0.05,\n",
       " 'sensitivity_interval': (2.567683475272766, 7.648542718465433),\n",
       " 'RV': 0.8172124624252319}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_level=0.05\n",
    "t_ind = 0\n",
    "\n",
    "theta, sigma, nu, sig = dr_sim(\n",
    "    RandomForestClassifier(), \n",
    "    RandomForestRegressor(), \n",
    "    2, \n",
    "    dpg_dict['y'],\n",
    "    dpg_dict['d'].astype(int),\n",
    "    dpg_dict['x']\n",
    ")\n",
    "\n",
    "result_dict = {\n",
    "    'alpha': sig_level,\n",
    "    'sensitivity_interval': sensitivity_interval(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level, 0.6, 0.6, 1),\n",
    "    'RV': RV(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level)\n",
    "}\n",
    "result_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 401k data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data = dml.datasets.fetch_401K()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "model selection for propensity model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RandomForestClassifier selected with score: 0.6435\n",
      "401k data treatment model selection:\n",
      "Best classifier: RandomForestClassifier(max_depth=10, n_estimators=200, random_state=42)\n",
      "Best score: 0.6435\n",
      "best rf score: 0.6435\n",
      "best lr score: 0.6578\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import GridSearchCV\n",
    "\n",
    "# Define parameter grid for model selection\n",
    "# Define parameter grid for model selection\n",
    "param_grid = {\n",
    "    'C': [0.1, 1, 10, 100]  # For LogisticRegression\n",
    "}\n",
    "\n",
    "# First find the best LogisticRegression model\n",
    "lr_grid_search = GridSearchCV(\n",
    "    estimator=LogisticRegression(max_iter=1000),\n",
    "    param_grid=param_grid,\n",
    "    cv=3,\n",
    "    scoring='accuracy'\n",
    ")\n",
    "lr_grid_search.fit(dml_data.x, dml_data.d.astype(int))\n",
    "best_lr = lr_grid_search.best_estimator_\n",
    "best_lr_score = lr_grid_search.best_score_\n",
    "\n",
    "# Find the best RandomForestClassifier\n",
    "rf_param_grid = {\n",
    "    'n_estimators': [50, 100, 200],\n",
    "    'max_depth': [None, 10, 20]\n",
    "}\n",
    "rf_grid_search = GridSearchCV(\n",
    "    estimator=RandomForestClassifier(random_state=42),\n",
    "    param_grid=rf_param_grid,\n",
    "    cv=3,\n",
    "    scoring='accuracy'\n",
    ")\n",
    "rf_grid_search.fit(dml_data.x, dml_data.d.astype(int))\n",
    "best_rf = rf_grid_search.best_estimator_\n",
    "best_rf_score = rf_grid_search.best_score_\n",
    "\n",
    "# Choose the overall best model\n",
    "if best_lr_score >= best_rf_score:\n",
    "    grid_search = lr_grid_search\n",
    "    print(f\"LogisticRegression selected with score: {best_lr_score:.4f}\")\n",
    "else:\n",
    "    grid_search = rf_grid_search\n",
    "print(f\"RandomForestClassifier selected with score: {best_rf_score:.4f}\")\n",
    "\n",
    "# Use the best estimator from grid search\n",
    "best_classifier = grid_search.best_estimator_\n",
    "print('401k data treatment model selection:')\n",
    "print(f\"Best classifier: {best_classifier}\")\n",
    "print(f\"Best score: {grid_search.best_score_:.4f}\")\n",
    "print(f\"best rf score: {best_rf_score:.4f}\")\n",
    "print(f\"best lr score: {best_lr_score:.4f}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== DoubleMLIRM Object ==================\n",
      "\n",
      "------------------ Data summary      ------------------\n",
      "Outcome variable: inuidur1\n",
      "Treatment variable(s): ['tg']\n",
      "Covariates: ['female', 'black', 'othrace', 'dep1', 'dep2', 'q2', 'q3', 'q4', 'q5', 'q6', 'agelt35', 'agegt54', 'durable', 'lusd', 'husd']\n",
      "Instrument variable(s): None\n",
      "No. Observations: 5099\n",
      "\n",
      "------------------ Score & algorithm ------------------\n",
      "Score function: ATE\n",
      "\n",
      "------------------ Machine learner   ------------------\n",
      "Learner ml_g: RandomForestRegressor()\n",
      "Learner ml_m: RandomForestClassifier()\n",
      "Out-of-sample Performance:\n",
      "Regression:\n",
      "Learner ml_g0 RMSE: [[1.2687313]]\n",
      "Learner ml_g1 RMSE: [[1.30580821]]\n",
      "Classification:\n",
      "Learner ml_m Log Loss: [[0.73043429]]\n",
      "\n",
      "------------------ Resampling        ------------------\n",
      "No. folds: 5\n",
      "No. repeated sample splits: 1\n",
      "\n",
      "------------------ Fit summary       ------------------\n",
      "        coef   std err         t     P>|t|     2.5 %    97.5 %\n",
      "tg -0.080843  0.215406 -0.375307  0.707432 -0.503031  0.341344\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Propensity predictions from learner RandomForestClassifier() for ml_m are close to zero or one (eps=1e-12).\n"
     ]
    }
   ],
   "source": [
    "dml_obj = dml.DoubleMLIRM(dml_data,\n",
    "                          ml_g=RandomForestRegressor(),\n",
    "                          ml_m=RandomForestClassifier(),\n",
    "                        #   ml_m=best_classifier,\n",
    "                          n_folds=5,\n",
    "                          score='ATE')\n",
    "dml_obj.fit()\n",
    "print(dml_obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "sensitivity_elements sigma2 and nu2 have to be positive. Got sigma2 [[[1.64234643]]] and nu2 [[[-103.64744455]]]. Most likely this is due to low quality learners (especially propensity scores).",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[82], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mdml_obj\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msensitivity_analysis\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcf_y\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.04\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcf_d\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.04\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrho\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1.\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(dml_obj\u001b[38;5;241m.\u001b[39msensitivity_summary)\n",
      "File \u001b[1;32mc:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\double_ml.py:1419\u001b[0m, in \u001b[0;36mDoubleML.sensitivity_analysis\u001b[1;34m(self, cf_y, cf_d, rho, level, null_hypothesis)\u001b[0m\n\u001b[0;32m   1417\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_framework \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   1418\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mApply fit() before sensitivity_analysis().\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m-> 1419\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_framework\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msensitivity_analysis\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1420\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcf_y\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcf_y\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1421\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcf_d\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcf_d\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1422\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrho\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrho\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1423\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlevel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlevel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1424\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnull_hypothesis\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnull_hypothesis\u001b[49m\n\u001b[0;32m   1425\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1427\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n",
      "File \u001b[1;32mc:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\double_ml_framework.py:613\u001b[0m, in \u001b[0;36mDoubleMLFramework.sensitivity_analysis\u001b[1;34m(self, cf_y, cf_d, rho, level, null_hypothesis)\u001b[0m\n\u001b[0;32m    609\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnull_hypothesis has to be of type float or np.ndarry. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    610\u001b[0m                     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mstr\u001b[39m(null_hypothesis)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m of type \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28mtype\u001b[39m(null_hypothesis))\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m was passed.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    612\u001b[0m \u001b[38;5;66;03m# compute sensitivity analysis\u001b[39;00m\n\u001b[1;32m--> 613\u001b[0m sensitivity_dict \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_calc_sensitivity_analysis\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcf_y\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcf_y\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcf_d\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcf_d\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mrho\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrho\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlevel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlevel\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    615\u001b[0m \u001b[38;5;66;03m# compute robustess values with respect to null_hypothesis\u001b[39;00m\n\u001b[0;32m    616\u001b[0m rv \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mfull(shape\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_n_thetas, fill_value\u001b[38;5;241m=\u001b[39mnp\u001b[38;5;241m.\u001b[39mnan)\n",
      "File \u001b[1;32mc:\\Users\\fabiovera\\AppData\\Local\\anaconda3\\envs\\dev_env2\\lib\\site-packages\\doubleml\\double_ml_framework.py:462\u001b[0m, in \u001b[0;36mDoubleMLFramework._calc_sensitivity_analysis\u001b[1;34m(self, cf_y, cf_d, rho, level)\u001b[0m\n\u001b[0;32m    459\u001b[0m psi_scaled \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_scaled_psi\n\u001b[0;32m    461\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (np\u001b[38;5;241m.\u001b[39many(sigma2 \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)) \u001b[38;5;241m|\u001b[39m (np\u001b[38;5;241m.\u001b[39many(nu2 \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)):\n\u001b[1;32m--> 462\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msensitivity_elements sigma2 and nu2 have to be positive. \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    463\u001b[0m                      \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGot sigma2 \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mstr\u001b[39m(sigma2)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m and nu2 \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mstr\u001b[39m(nu2)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    464\u001b[0m                      \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mMost likely this is due to low quality learners (especially propensity scores).\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m    466\u001b[0m \u001b[38;5;66;03m# elementwise operations\u001b[39;00m\n\u001b[0;32m    467\u001b[0m confounding_strength \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mmultiply(np\u001b[38;5;241m.\u001b[39mabs(rho), np\u001b[38;5;241m.\u001b[39msqrt(np\u001b[38;5;241m.\u001b[39mmultiply(cf_y, np\u001b[38;5;241m.\u001b[39mdivide(cf_d, \u001b[38;5;241m1.0\u001b[39m\u001b[38;5;241m-\u001b[39mcf_d))))\n",
      "\u001b[1;31mValueError\u001b[0m: sensitivity_elements sigma2 and nu2 have to be positive. Got sigma2 [[[1.64234643]]] and nu2 [[[-103.64744455]]]. Most likely this is due to low quality learners (especially propensity scores)."
     ]
    }
   ],
   "source": [
    "dml_obj.sensitivity_analysis(cf_y=0.04, cf_d=0.04, rho=1.)\n",
    "print(dml_obj.sensitivity_summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'alpha': 0.05,\n",
       " 'sensitivity_interval': (-0.15765123322817087, 0.01212238369880464),\n",
       " 'RV': 0}"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_level=0.05\n",
    "t_ind = 0\n",
    "\n",
    "theta, sigma, nu, sig = dr_sim(\n",
    "    best_classifier, \n",
    "    RandomForestRegressor(), \n",
    "    5, \n",
    "    dml_data.y,\n",
    "    dml_data.d.astype(int),\n",
    "    dml_data.x\n",
    ")\n",
    "\n",
    "result_dict = {\n",
    "    'alpha': sig_level,\n",
    "    'sensitivity_interval': sensitivity_interval(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level, 0.00001, 0.0001, 1),\n",
    "    'RV': RV(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level)\n",
    "}\n",
    "result_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### bonus data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "dml_data = dml.datasets.fetch_bonus()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "model selection for propensity model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LogisticRegression selected with score: -0.6442\n",
      "RandomForestClassifier selected with score: -0.6624\n",
      "bonus data treatment model selection:\n",
      "Best classifier: LogisticRegression(C=0.1, max_iter=1000)\n",
      "Best score: -0.6442\n",
      "best rf score: -0.6624\n",
      "best lr score: -0.6442\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import GridSearchCV\n",
    "\n",
    "# Define parameter grid for model selection\n",
    "# Define parameter grid for model selection\n",
    "param_grid = {\n",
    "    'C': [0.1, 1, 10, 100]  # For LogisticRegression\n",
    "}\n",
    "\n",
    "# First find the best LogisticRegression model\n",
    "lr_grid_search = GridSearchCV(\n",
    "    estimator=LogisticRegression(max_iter=1000),\n",
    "    param_grid=param_grid,\n",
    "    cv=3,\n",
    "    scoring='neg_log_loss'\n",
    ")\n",
    "lr_grid_search.fit(dml_data.x, dml_data.d.astype(int))\n",
    "best_lr = lr_grid_search.best_estimator_\n",
    "best_lr_score = lr_grid_search.best_score_\n",
    "\n",
    "# Find the best RandomForestClassifier\n",
    "rf_param_grid = {\n",
    "    'n_estimators': [50, 100, 200],\n",
    "    'max_depth': [None, 10, 20]\n",
    "}\n",
    "rf_grid_search = GridSearchCV(\n",
    "    estimator=RandomForestClassifier(random_state=42),\n",
    "    param_grid=rf_param_grid,\n",
    "    cv=3,\n",
    "    scoring='neg_log_loss'\n",
    ")\n",
    "rf_grid_search.fit(dml_data.x, dml_data.d.astype(int))\n",
    "best_rf = rf_grid_search.best_estimator_\n",
    "best_rf_score = rf_grid_search.best_score_\n",
    "\n",
    "# Choose the overall best model\n",
    "if best_lr_score >= best_rf_score:\n",
    "    grid_search = lr_grid_search\n",
    "    print(f\"LogisticRegression selected with score: {best_lr_score:.4f}\")\n",
    "else:\n",
    "    grid_search = rf_grid_search\n",
    "print(f\"RandomForestClassifier selected with score: {best_rf_score:.4f}\")\n",
    "\n",
    "# Use the best estimator from grid search\n",
    "best_classifier = grid_search.best_estimator_\n",
    "print('bonus data treatment model selection:')\n",
    "print(f\"Best classifier: {best_classifier}\")\n",
    "print(f\"Best score: {grid_search.best_score_:.4f}\")\n",
    "print(f\"best rf score: {best_rf_score:.4f}\")\n",
    "print(f\"best lr score: {best_lr_score:.4f}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== DoubleMLIRM Object ==================\n",
      "\n",
      "------------------ Data summary      ------------------\n",
      "Outcome variable: inuidur1\n",
      "Treatment variable(s): ['tg']\n",
      "Covariates: ['female', 'black', 'othrace', 'dep1', 'dep2', 'q2', 'q3', 'q4', 'q5', 'q6', 'agelt35', 'agegt54', 'durable', 'lusd', 'husd']\n",
      "Instrument variable(s): None\n",
      "No. Observations: 5099\n",
      "\n",
      "------------------ Score & algorithm ------------------\n",
      "Score function: ATE\n",
      "\n",
      "------------------ Machine learner   ------------------\n",
      "Learner ml_g: RandomForestRegressor()\n",
      "Learner ml_m: LogisticRegression(C=0.1, max_iter=1000)\n",
      "Out-of-sample Performance:\n",
      "Regression:\n",
      "Learner ml_g0 RMSE: [[1.29747417]]\n",
      "Learner ml_g1 RMSE: [[1.33043293]]\n",
      "Classification:\n",
      "Learner ml_m Log Loss: [[0.64522511]]\n",
      "\n",
      "------------------ Resampling        ------------------\n",
      "No. folds: 2\n",
      "No. repeated sample splits: 1\n",
      "\n",
      "------------------ Fit summary       ------------------\n",
      "        coef   std err         t     P>|t|     2.5 %    97.5 %\n",
      "tg -0.059058  0.038039 -1.552562  0.120528 -0.133613  0.015497\n"
     ]
    }
   ],
   "source": [
    "dml_obj = dml.DoubleMLIRM(dml_data,\n",
    "                          ml_g=RandomForestRegressor(),\n",
    "                          ml_m=best_classifier,\n",
    "                          n_folds=2,\n",
    "                          score='ATE',)\n",
    "dml_obj.fit()\n",
    "print(dml_obj)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================== Sensitivity Analysis ==================\n",
      "\n",
      "------------------ Scenario          ------------------\n",
      "Significance Level: level=0.95\n",
      "Sensitivity parameters: cf_y=0.0; cf_d=0.0, rho=1.0\n",
      "\n",
      "------------------ Bounds with CI    ------------------\n",
      "   CI lower  theta lower     theta  theta upper  CI upper\n",
      "0 -0.121627    -0.059058 -0.059058    -0.059058  0.003511\n",
      "\n",
      "------------------ Robustness Values ------------------\n",
      "   H_0    RV (%)   RVa (%)\n",
      "0  0.0  2.128758  0.000434\n"
     ]
    }
   ],
   "source": [
    "dml_obj.sensitivity_analysis(cf_y=0.04, cf_d=0.04, rho=1.)\n",
    "print(dml_obj.sensitivity_summary)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'alpha': 0.05,\n",
       " 'sensitivity_interval': (-0.29151616548820364, 0.1365988298703495),\n",
       " 'RV': 0}"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sig_level=0.05\n",
    "t_ind = 0\n",
    "\n",
    "theta, sigma, nu, sig = dr_sim(\n",
    "    best_classifier, \n",
    "    RandomForestRegressor(), \n",
    "    5, \n",
    "    dml_data.y,\n",
    "    dml_data.d.astype(int),\n",
    "    dml_data.x\n",
    ")\n",
    "\n",
    "result_dict = {\n",
    "    'alpha': sig_level,\n",
    "    'sensitivity_interval': sensitivity_interval(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level, 0.05, 0.05, 1),\n",
    "    'RV': RV(theta[t_ind], sigma[t_ind], nu[t_ind], sig[t_ind], sig_level)\n",
    "}\n",
    "result_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dev_env2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
